#!/usr/bin/env python3

import rosbag
import matplotlib.pyplot as plt
import numpy as np

# 定义函数提取 XYZ 数据和时间戳
def extract_xyz_from_bag(bag_file, odom_topic, pose_topic):
    # 初始化存储数据的列表
    odom_time, odom_x, odom_y, odom_z = [], [], [], []
    pose_time, pose_x, pose_y, pose_z = [], [], [], []

    # 打开 bag 文件
    with rosbag.Bag(bag_file, 'r') as bag:
        # 获取 bag 文件的起始时间
        start_time = bag.get_start_time()

        # 遍历 bag 文件中的消息
        for topic, msg, t in bag.read_messages(topics=[odom_topic, pose_topic]):
            # 计算相对于起始时间的时间戳（秒）
            timestamp = (t.to_sec() - start_time)

            if topic == odom_topic:
                # 提取 /baton/stereo3/odometry 的 XYZ 和时间戳
                odom_time.append(timestamp)
                odom_x.append(msg.pose.pose.position.x)
                odom_y.append(msg.pose.pose.position.y)
                odom_z.append(msg.pose.pose.position.z)
            elif topic == pose_topic:
                # 提取 /vrpn_client_node/uav1/pose 的 XYZ 和时间戳
                pose_time.append(timestamp)
                pose_x.append(msg.pose.position.x)
                pose_y.append(msg.pose.position.y)
                pose_z.append(msg.pose.position.z)

    return odom_time, odom_x, odom_y, odom_z, pose_time, pose_x, pose_y, pose_z

# 定义函数绘制 XYZ 数据
def plot_xyz(odom_time, odom_x, odom_y, odom_z, pose_time, pose_x, pose_y, pose_z):
    # 创建图形
    plt.figure(figsize=(12, 8))

    # 绘制 /baton/stereo3/odometry 的 XYZ
    plt.plot(odom_time, odom_x, label='/baton/stereo3/odometry X', color='r', linestyle='-')
    plt.plot(odom_time, odom_y, label='/baton/stereo3/odometry Y', color='g', linestyle='-')
    plt.plot(odom_time, odom_z, label='/baton/stereo3/odometry Z', color='b', linestyle='-')

    # 绘制 /vrpn_client_node/uav1/pose 的 XYZ
    plt.plot(pose_time, pose_x, label='/vrpn_client_node/uav1/pose X', color='r', linestyle='--')
    plt.plot(pose_time, pose_y, label='/vrpn_client_node/uav1/pose Y', color='g', linestyle='--')
    plt.plot(pose_time, pose_z, label='/vrpn_client_node/uav1/pose Z', color='b', linestyle='--')

    # 设置图形属性
    plt.xlabel('Time (s)')
    plt.ylabel('Position (m)')
    plt.title('XYZ Position over Time')
    plt.legend()
    plt.grid(True)

    # 设置 X 轴刻度为 1000ms (1s)
    min_time = min(min(odom_time), min(pose_time))  # 最小时间
    max_time = max(max(odom_time), max(pose_time))  # 最大时间
    plt.xticks(np.arange(min_time, max_time + 1, 1), rotation=45)  # 1000ms 间隔

    # 调整布局
    plt.tight_layout()
    plt.show()

# 主函数
if __name__ == "__main__":
    # 定义 bag 文件路径和话题名称
    bag_file = 'wheel_circle.bag'
    odom_topic = '/sunray/robobaton_mini/odom'
    pose_topic = '/vrpn_client_node/uav1/pose'

    # 提取 XYZ 数据和时间戳
    odom_time, odom_x, odom_y, odom_z, pose_time, pose_x, pose_y, pose_z = extract_xyz_from_bag(bag_file, odom_topic, pose_topic)

    # 绘制 XYZ 数据
    plot_xyz(odom_time, odom_x, odom_y, odom_z, pose_time, pose_x, pose_y, pose_z)
